{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7737092",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "\n",
    "origin_data_path = 'path/to/qasper'\n",
    "\n",
    "\n",
    "def process(doc):\n",
    "    seg_point_count = [0]*10\n",
    "    seg_point_count[0] = 1\n",
    "    \n",
    "    doc_text = doc.get('title', '')\n",
    "    doc_text = doc_text + '\\n' + doc.get('abstract', '')\n",
    "    for sec in doc['full_text']:\n",
    "        if sec.get('section_name', None):\n",
    "            titles = [t.strip() for t in sec['section_name'].split(':::')]\n",
    "            paragraphs = sec['paragraphs']\n",
    "            level = len(titles)\n",
    "            title_text = f'{\"#\"*level} ' + titles[-1]\n",
    "            section_text = '\\n'.join(paragraphs)\n",
    "            doc_text = doc_text + '\\n' + title_text + '\\n' + section_text\n",
    "            seg_point_count[level] += 1\n",
    "        else:\n",
    "            paragraphs = sec['paragraphs']\n",
    "            section_text = '\\n'.join(paragraphs)\n",
    "            doc_text = doc_text + '\\n' + section_text\n",
    "        \n",
    "    \n",
    "    return doc_text, seg_point_count\n",
    "\n",
    "    \n",
    "splits = ['train', 'dev', 'test']\n",
    "for split in splits:\n",
    "    seg_point_count = [0]*10\n",
    "    save_path = f'corpus/qasper/{split}_file'\n",
    "    os.makedirs(save_path, exist_ok=True)\n",
    "    \n",
    "    input_path = f'{origin_data_path}/qasper-{split}-v0.3.json'\n",
    "    doc_dict = json.load(open(input_path))\n",
    "    for _id in doc_dict:\n",
    "        doc_text, spc = process(doc_dict[_id])\n",
    "        with open(f'{save_path}/{_id}.txt', 'w') as f:\n",
    "            f.write(doc_text)\n",
    "        for i in range(10):\n",
    "            seg_point_count[i] += spc[i]\n",
    "            \n",
    "    print(split, seg_point_count)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35a65816",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "\n",
    "origin_data_path = 'path/to/gov-report'\n",
    "\n",
    "\n",
    "def process(section, level, spc):\n",
    "    if section['section_title'].strip()==\"\" and len(section['paragraphs']) == 0:\n",
    "        return '', False\n",
    "    res = '#'*level + ' ' + (f\"{section['section_title']}\\n\" if ('section_title' in section and section['section_title'].strip() != '') else '') + '\\n'.join(section['paragraphs']) + '\\n'\n",
    "    spc[level] += 1\n",
    "    for subs in section['subsections']:\n",
    "        _, flag = process(subs, level+1, spc)\n",
    "        if not flag:\n",
    "            return '', False\n",
    "        else:\n",
    "            res += _\n",
    "\n",
    "    return res, True\n",
    "\n",
    "    \n",
    "splits = ['train', 'valid', 'test']\n",
    "save_splits_map = {\n",
    "    'train': 'train',\n",
    "    'valid': 'dev',\n",
    "    'test': 'test'\n",
    "}\n",
    "for split in splits:\n",
    "    seg_point_count = [0]*10\n",
    "    save_path = f'corpus/gov-report_5w/{save_splits_map[split]}_file'\n",
    "    os.makedirs(save_path, exist_ok=True)\n",
    "    \n",
    "    for n in tqdm(open(f'{origin_data_path}/split_ids/gao_{split}.ids', 'r').readlines()):\n",
    "        n = n.strip()\n",
    "        d = json.load(open(f'{origin_data_path}/gao/{n}.json', 'r'))\n",
    "        doc = ''\n",
    "        spc = [0]*10\n",
    "        spc[0] += 1\n",
    "        for subs in d['report']:\n",
    "            _, flag = process(subs, 1, spc)\n",
    "            if not flag:\n",
    "                doc = ''\n",
    "                break\n",
    "            else:\n",
    "                doc += _\n",
    "        if doc != '' and len(doc) >= 50000:\n",
    "            for i in range(10):\n",
    "                seg_point_count[i] += spc[i]\n",
    "            with open(f'{save_path}/gao_{n}.txt', 'w') as f:\n",
    "                f.write(doc)\n",
    "    \n",
    "    \n",
    "    for n in tqdm(open(f'{origin_data_path}/split_ids/crs_{split}.ids', 'r').readlines()):\n",
    "        n = n.strip()\n",
    "        spc = [0]*10\n",
    "        spc[0] += 1\n",
    "        d = json.load(open(f'{origin_data_path}/crs/{n}.json', 'r'))\n",
    "        doc, flag = process(d['reports'], 1, spc)\n",
    "        if doc != '' and len(doc) >= 50000:\n",
    "            for i in range(10):\n",
    "                seg_point_count[i] += spc[i]\n",
    "            with open(f'{save_path}/crs_{n}.txt', 'w') as f:\n",
    "                f.write(doc)\n",
    "    \n",
    "    print(split, seg_point_count)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "187d9e45",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "origin_data_path = 'path/to/wiki_727'\n",
    "splits = ['train', 'dev', 'test']\n",
    "for split in splits:\n",
    "    seg_point_count = [0]*10\n",
    "    save_path = f'corpus/wiki_727_5w/{split}_file'\n",
    "    os.makedirs(save_path, exist_ok=True)\n",
    "    \n",
    "    for root, dirnames, filenames in tqdm(os.walk(f'{origin_data_path}/{split}/')):\n",
    "        for fn in filenames:\n",
    "            with open(os.path.join(root, fn), 'r', encoding='utf-8') as f:\n",
    "                data = f.read()\n",
    "            if len(data) < 50000:\n",
    "                continue\n",
    "                \n",
    "            spc = [0]*10\n",
    "            save_name = root.lstrip(f'{origin_data_path}/{split}/').replace('/', '_') + '_' + fn\n",
    "            with open(os.path.join(root, fn), 'r', encoding='utf-8') as f:\n",
    "                lines = f.readlines()\n",
    "            spc[0] += 1\n",
    "            for i, l in enumerate(lines):\n",
    "                if l.startswith('========,'):\n",
    "                    level = int(l.split(',')[1].strip())\n",
    "                    spc[level] += 1\n",
    "    \n",
    "            if spc[1] <= 1:\n",
    "                for i, l in enumerate(lines):\n",
    "                    if l.startswith('========,'):\n",
    "                        level = int(l.split(',')[1].strip())\n",
    "                        lines[i] = f\"{'#' * (level if level == 1 else level - 1)} {l.split(',')[2]}\"\n",
    "                spc.pop(1)\n",
    "                spc[1] += 1\n",
    "                spc.append(0)\n",
    "            else:\n",
    "                for i, l in enumerate(lines):\n",
    "                    if l.startswith('========,'):\n",
    "                        level = int(l.split(',')[1].strip())\n",
    "                        lines[i] = f\"{'#' * level} {l.split(',')[2]}\"\n",
    "            \n",
    "            for i in range(10):\n",
    "                seg_point_count[i] += spc[i]\n",
    "            with open(f\"{save_path}/{save_name}\", 'w', encoding='utf-8') as f:\n",
    "                f.writelines(lines)\n",
    "    \n",
    "    print(split, seg_point_count)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64bcdd4b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
